package main

import (
	"context"
	"encoding/json"
	"github.com/360EntSecGroup-Skylar/excelize"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"log"
	"os"
	"strconv"
	"sync"
)

func main() {
	// 获取配置信息
	config, err := initConfig()
	if err != nil {
		log.Println("config.json 文件找不到或出现异常，请检查文件目录和内容是否正确")
	}

	// 连接 MongoDB
	clientOptions := options.Client().ApplyURI(config.URI)
	client, err := mongo.Connect(context.TODO(), clientOptions)
	if err != nil {
		panic(err.Error())
	}

	// 程序结束后断开 MongoDB 连接
	defer func() {
		if err = client.Disconnect(context.TODO()); err != nil {
			panic(err.Error())
		}
	}()

	// 选择要连接哪个库
	database := client.Database(config.Database)

	log.Println("开始导入")

	wg := sync.WaitGroup{}
	wg.Add(5)

	go loadStudentData(database, &wg)
	go loadTeacherData(database, &wg)
	go loadCourseData(database, &wg)
	go loadStudentCourseData(database, &wg)
	go loadTeacherCourseData(database, &wg)

	wg.Wait()

	log.Println("导入完毕")
}

// Config 配置信息
type Config struct {
	URI      string `json:"uri"`
	Database string `json:"database"` // 数据库名
}

// 初始化配置
func initConfig() (*Config, error) {
	// 读取 config.json 中的文件
	file, err := os.Open("./config.json")
	if err != nil {
		return nil, err
	}

	// 解码
	config := Config{}
	err = json.NewDecoder(file).Decode(&config)
	if err != nil {
		return nil, err
	}

	return &config, nil
}

// 导入 student 数据
func loadStudentData(database *mongo.Database, wg *sync.WaitGroup) {
	// 获取 student 集合
	student := database.Collection("student")

	// 导入文件
	studentData, err := excelize.OpenFile("./student.xlsx")
	if err != nil {
		log.Println("student.xlsx 导入失败，请检查文件目录是否正确")
		return
	}

	// 获取 student 数据
	rows := studentData.GetRows("Sheet1")

	// 创建数组
	length := len(rows)
	docs := make([]interface{}, 0, length-1)

	// 根据表格中第一行的值，建立列名与索引的对应关系
	studentMap := map[string]int{}
	for i := 0; i < 7; i++ {
		studentMap[rows[0][i]] = i
	}

	// 获取每一列的索引
	sidIndex := studentMap["SID"]
	nameIndex := studentMap["NAME"]
	sexIndex := studentMap["SEX"]
	ageIndex := studentMap["AGE"]
	dnameIndex := studentMap["DNAME"]
	classIndex := studentMap["CLASS"]
	birthdayIndex := studentMap["BIRTHDAY"]

	// 遍历，进行数据格式化
	for index, row := range rows {
		// 跳过列名所在的行
		if index == 0 {
			continue
		}

		// 将年龄转化为整型
		age, err := strconv.Atoi(row[ageIndex])
		if err != nil {
			log.Println("第", index, "行的年龄数据存在异常")
		}

		docs = append(docs, bson.M{
			"sid":      row[sidIndex],
			"name":     row[nameIndex],
			"sex":      row[sexIndex],
			"age":      age,
			"dname":    row[dnameIndex],
			"class":    row[classIndex],
			"birthday": row[birthdayIndex],
		})
	}

	result, err := student.InsertMany(context.TODO(), docs)
	if err != nil {
		log.Println("student 插入数据失败")
		return
	}

	log.Println("student 成功插入", len(result.InsertedIDs), "行数据")

	// 解锁
	wg.Done()
}

// 导入 teacher 数据
func loadTeacherData(database *mongo.Database, wg *sync.WaitGroup) {
	teacher := database.Collection("teacher")

	teacherData, err := excelize.OpenFile("./teacher.xlsx")
	if err != nil {
		log.Println("teacher.xlsx 导入失败，请检查文件目录是否正确")
		return
	}

	rows := teacherData.GetRows("Sheet1")

	length := len(rows)
	docs := make([]interface{}, 0, length-1)

	teacherMap := map[string]int{}
	for i := 0; i < 5; i++ {
		teacherMap[rows[0][i]] = i
	}

	tidIndex := teacherMap["TID"]
	nameIndex := teacherMap["NAME"]
	sexIndex := teacherMap["SEX"]
	ageIndex := teacherMap["AGE"]
	dnameIndex := teacherMap["DNAME"]

	for index, row := range rows {
		if index == 0 {
			continue
		}

		age, err := strconv.Atoi(row[ageIndex])
		if err != nil {
			log.Println("第", index, "行的年龄数据存在异常")
		}

		docs = append(docs, bson.M{
			"tid":   row[tidIndex],
			"name":  row[nameIndex],
			"sex":   row[sexIndex],
			"age":   age,
			"dname": row[dnameIndex],
		})
	}

	result, err := teacher.InsertMany(context.TODO(), docs)
	if err != nil {
		log.Println("teacher 插入数据失败")
		return
	}

	log.Println("teacher 成功插入", len(result.InsertedIDs), "行数据")

	wg.Done()
}

// 导入 course 数据
func loadCourseData(database *mongo.Database, wg *sync.WaitGroup) {
	course := database.Collection("course")

	courseData, err := excelize.OpenFile("./course.xlsx")
	if err != nil {
		log.Println("course.xlsx 导入失败，请检查文件目录是否正确")
		return
	}

	rows := courseData.GetRows("Sheet1")

	length := len(rows)
	docs := make([]interface{}, 0, length-1)

	courseMap := map[string]int{}
	for i := 0; i < 4; i++ {
		courseMap[rows[0][i]] = i
	}

	cidIndex := courseMap["CID"]
	nameIndex := courseMap["NAME"]
	fcidIndex := courseMap["FCID"]
	creditIndex := courseMap["CREDIT"]

	for index, row := range rows {
		if index == 0 {
			continue
		}

		credit, err := strconv.Atoi(row[creditIndex])
		if err != nil {
			log.Println("第", index, "行的学分数据存在异常")
		}

		docs = append(docs, bson.M{
			"cid":    row[cidIndex],
			"name":   row[nameIndex],
			"fcid":   row[fcidIndex],
			"credit": credit,
		})
	}

	result, err := course.InsertMany(context.TODO(), docs)
	if err != nil {
		log.Println("course 插入数据失败")
		return
	}

	log.Println("course 成功插入", len(result.InsertedIDs), "行数据")

	wg.Done()
}

// 导入 student_course 数据
func loadStudentCourseData(database *mongo.Database, wg *sync.WaitGroup) {
	studentCourse := database.Collection("student_course")

	studentCourseData, err := excelize.OpenFile("./student_course.xlsx")
	if err != nil {
		log.Println("student_course.xlsx 导入失败，请检查文件目录是否正确")
		return
	}

	rows := studentCourseData.GetRows("Sheet1")

	length := len(rows)
	docs := make([]interface{}, 0, length-1)

	studentCourseMap := map[string]int{}
	for i := 0; i < 4; i++ {
		studentCourseMap[rows[0][i]] = i
	}

	sidIndex := studentCourseMap["SID"]
	cidIndex := studentCourseMap["CID"]
	scoreIndex := studentCourseMap["SCORE"]
	tidIndex := studentCourseMap["TID"]

	for index, row := range rows {
		if index == 0 {
			continue
		}

		score, err := strconv.Atoi(row[scoreIndex])
		if err != nil {
			log.Println("第", index, "行的成绩数据存在异常")
		}

		docs = append(docs, bson.M{
			"sid":   row[sidIndex],
			"cid":   row[cidIndex],
			"score": score,
			"tid":   row[tidIndex],
		})
	}

	result, err := studentCourse.InsertMany(context.TODO(), docs)
	if err != nil {
		log.Println("student_course 插入数据失败")
		return
	}

	log.Println("student_course 成功插入", len(result.InsertedIDs), "行数据")

	wg.Done()
}

// 导入 teacher_course 数据
func loadTeacherCourseData(database *mongo.Database, wg *sync.WaitGroup) {
	teacherCourse := database.Collection("teacher_course")

	teacherCourseData, err := excelize.OpenFile("./teacher_course.xlsx")
	if err != nil {
		log.Println("teacher_course.xlsx 导入失败，请检查文件目录是否正确")
		return
	}

	rows := teacherCourseData.GetRows("Sheet1")

	length := len(rows)
	docs := make([]interface{}, 0, length-1)

	teacherCourseMap := map[string]int{}
	for i := 0; i < 2; i++ {
		teacherCourseMap[rows[0][i]] = i
	}

	cidIndex := teacherCourseMap["CID"]
	tidIndex := teacherCourseMap["TID"]

	for index, row := range rows {
		if index == 0 {
			continue
		}

		docs = append(docs, bson.M{
			"cid": row[cidIndex],
			"tid": row[tidIndex],
		})
	}

	result, err := teacherCourse.InsertMany(context.TODO(), docs)
	if err != nil {
		log.Println("teacher_course 插入数据失败")
		return
	}

	log.Println("teacher_course 成功插入", len(result.InsertedIDs), "行数据")

	wg.Done()
}
